"""字节跳动豆包模型适配器实现。"""

from typing import AsyncGenerator, Dict, Any, List
import aiohttp
import json
import time
import hashlib

from .base import BaseModelAdapter, ModelConfig, ModelResponse


class DoubaoAdapter(BaseModelAdapter):
    """豆包（字节跳动）大模型适配器。"""
    
    # 定价（每千tokens，人民币）
    PRICING = {
        "skylark-chat": {"prompt": 0.003, "completion": 0.009},
        "skylark-lite": {"prompt": 0.0008, "completion": 0.0025},
        "skylark-plus": {"prompt": 0.005, "completion": 0.015},
        "skylark-pro": {"prompt": 0.008, "completion": 0.025}
    }
    
    def __init__(self, config: ModelConfig):
        """初始化豆包适配器。
        
        需要在extra_params中提供:
        - endpoint_id: 模型端点ID
        """
        super().__init__(config)
        self.endpoint_id = config.extra_params.get("endpoint_id", "")
        
    def _initialize_client(self) -> None:
        """初始化豆包客户端。"""
        # 火山引擎API地址
        self.base_url = self.config.api_base or "https://ark.cn-beijing.volces.com/api/v3"
    
    async def generate(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> ModelResponse:
        """生成响应。
        
        Args:
            messages: 对话消息
            **kwargs: 额外参数
            
        Returns:
            模型响应
        """
        url = f"{self.base_url}/chat/completions"
        
        # 生成签名
        timestamp = str(int(time.time()))
        sign = self._generate_signature(timestamp)
        
        headers = {
            "Authorization": f"Bearer {self.config.api_key}",
            "Content-Type": "application/json",
            "X-Timestamp": timestamp,
            "X-Signature": sign
        }
        
        # 构建请求数据
        data = {
            "model": self.endpoint_id or self.config.model_name,
            "messages": self._convert_messages(messages),
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens,
            "top_p": self.config.top_p,
            "frequency_penalty": self.config.frequency_penalty,
            "presence_penalty": self.config.presence_penalty,
            "stream": False
        }
        
        # 豆包特有参数
        if "plugins" in kwargs:
            data["plugins"] = kwargs.pop("plugins")  # 插件支持
        if "search" in kwargs:
            data["search"] = kwargs.pop("search")  # 搜索增强
        
        data.update(kwargs)
        data.update(self.config.extra_params)
        
        async with aiohttp.ClientSession() as session:
            async with session.post(
                url,
                headers=headers,
                json=data,
                timeout=aiohttp.ClientTimeout(total=self.config.timeout)
            ) as response:
                result = await response.json()
                
                if "error" in result:
                    raise Exception(f"豆包API错误: {result['error'].get('message', 'Unknown error')}")
                
                choice = result["choices"][0]
                usage = result.get("usage", {})
                
                return ModelResponse(
                    content=choice["message"]["content"],
                    model=result.get("model", self.config.model_name),
                    usage={
                        "prompt_tokens": usage.get("prompt_tokens", 0),
                        "completion_tokens": usage.get("completion_tokens", 0),
                        "total_tokens": usage.get("total_tokens", 0)
                    },
                    finish_reason=choice.get("finish_reason", "stop"),
                    metadata={
                        "id": result.get("id", ""),
                        "created": result.get("created", 0)
                    }
                )
    
    async def stream(
        self,
        messages: List[Dict[str, str]],
        **kwargs
    ) -> AsyncGenerator[str, None]:
        """流式生成响应。
        
        Args:
            messages: 对话消息
            **kwargs: 额外参数
            
        Yields:
            内容块
        """
        url = f"{self.base_url}/chat/completions"
        
        timestamp = str(int(time.time()))
        sign = self._generate_signature(timestamp)
        
        headers = {
            "Authorization": f"Bearer {self.config.api_key}",
            "Content-Type": "application/json",
            "X-Timestamp": timestamp,
            "X-Signature": sign
        }
        
        data = {
            "model": self.endpoint_id or self.config.model_name,
            "messages": self._convert_messages(messages),
            "temperature": self.config.temperature,
            "max_tokens": self.config.max_tokens,
            "top_p": self.config.top_p,
            "stream": True
        }
        
        data.update(kwargs)
        data.update(self.config.extra_params)
        
        async with aiohttp.ClientSession() as session:
            async with session.post(
                url,
                headers=headers,
                json=data,
                timeout=aiohttp.ClientTimeout(total=self.config.timeout)
            ) as response:
                async for line in response.content:
                    if line:
                        line_str = line.decode('utf-8').strip()
                        if line_str.startswith("data: "):
                            data_str = line_str[6:]
                            if data_str == "[DONE]":
                                break
                            
                            try:
                                chunk = json.loads(data_str)
                                if "choices" in chunk and chunk["choices"]:
                                    delta = chunk["choices"][0].get("delta", {})
                                    content = delta.get("content", "")
                                    if content:
                                        yield content
                            except json.JSONDecodeError:
                                continue
    
    def _convert_messages(self, messages: List[Dict[str, str]]) -> List[Dict[str, str]]:
        """转换消息格式以适配豆包API。
        
        Args:
            messages: 原始消息列表
            
        Returns:
            转换后的消息列表
        """
        converted = []
        for msg in messages:
            role = msg["role"]
            # 豆包使用的角色名称
            if role == "system":
                role = "system"
            elif role == "assistant":
                role = "assistant"
            else:
                role = "user"
            
            converted.append({
                "role": role,
                "content": msg["content"]
            })
        
        return converted
    
    def _generate_signature(self, timestamp: str) -> str:
        """生成API请求签名。
        
        Args:
            timestamp: 时间戳
            
        Returns:
            签名字符串
        """
        # 简单的签名实现，实际可能更复杂
        secret_key = self.config.extra_params.get("secret_key", "")
        sign_str = f"{self.config.api_key}{timestamp}{secret_key}"
        return hashlib.sha256(sign_str.encode()).hexdigest()
    
    def count_tokens(self, text: str) -> int:
        """计算文本的token数量。
        
        豆包的token计算方式。
        
        Args:
            text: 输入文本
            
        Returns:
            Token数量
        """
        # 中文字符
        chinese_chars = sum(1 for char in text if '\u4e00' <= char <= '\u9fff')
        # 英文单词
        english_words = len([w for w in text.split() if w.isascii()])
        
        # 豆包的估算：中文1字符约1 token，英文1单词约1.3 token
        return chinese_chars + int(english_words * 1.3)
    
    def estimate_cost(self, prompt_tokens: int, completion_tokens: int) -> float:
        """估算使用成本。
        
        Args:
            prompt_tokens: 提示词token数
            completion_tokens: 生成token数
            
        Returns:
            预估成本（美元）
        """
        # 根据模型名称判断定价
        model_key = "skylark-chat"  # 默认
        for key in self.PRICING:
            if key in self.config.model_name.lower():
                model_key = key
                break
        
        pricing = self.PRICING[model_key]
        prompt_cost = (prompt_tokens / 1000) * pricing["prompt"]
        completion_cost = (completion_tokens / 1000) * pricing["completion"]
        
        # 转换为美元
        cny_to_usd = 0.14
        return (prompt_cost + completion_cost) * cny_to_usd